"""
Simplified AES Encryption
"""

BINARY_VALUE = [
                "0000",
                "0001",
                "0010",
                "0011",
                "0100",
                "0101",
                "0110",
                "0111",
                "1000",
                "1001",
                "1010",
                "1011",
                "1100",
                "1101",
                "1110",
                "1111"
                ]

def key_expansion(key):
    """
    Key Expansion
    """
    RCON1 = "10000000"
    RCON2 = "00110000"
    w_zero = key[:8]
    w_one = key[8:]
    w_two = _xor(_xor(w_zero, RCON1), nib_sub_8(rot_nib(w_one)))
    w_three = _xor(w_two, w_one)
    w_four = _xor(_xor(w_two, RCON2), nib_sub_8(rot_nib(w_three)))
    w_five = _xor(w_four, w_three)
    return w_zero + w_one, w_two + w_three, w_four + w_five

def add_round_key(message, key):
    """
    Add Round Key
    """
    return _xor(message, key)

def s_box(binary_string):
    """
    AES S-Box
    """
    box = [
           ["1001", "0100", "1010", "1011"],
           ["1101", "0001", "1000", "0101"],
           ["0110", "0010", "0000", "0011"],
           ["1100", "1110", "1111", "0111"]
           ]
    row = int(binary_string[:2], 2)
    column = int(binary_string[2:], 2)
    return box[row][column]

def nib_sub_8(binary_string):
    """
    8-bit Nibble Substitution
    """
    return s_box(binary_string[:4]) + s_box(binary_string[4:])

def nib_sub_16(binary_string):
    """
    16-bit Nibble Substitution
    """
    return s_box(binary_string[:4]) + s_box(binary_string[4:8]) + \
                        s_box(binary_string[8:12]) + s_box(binary_string[12:16])

def rot_nib(binary_string):
    """
    Rotate Nibbles
    """
    return binary_string[4:] + binary_string[:4]

def shift_row(binary_string):
    """
    Shift Row
    """
    return binary_string[:4] + binary_string[12:16] + binary_string[8:12] + binary_string[4:8]

def mix_columns(binary_string):
    """
    Mix Columns
    """
    s = [
         [binary_string[:4], binary_string[8:12]],
         [binary_string[4:8], binary_string[12:16]]
         ]
    s_prime = [
               ["", ""],
               ["", ""]
               ]
    
    # Addition in GF(2^n) is equivalent to XOR
    s_prime[0][0] = _xor(s[0][0], _mult_gf_24("0100", s[1][0]))
    s_prime[1][0] = _xor(_mult_gf_24("0100", s[0][0]), s[1][0])
    s_prime[0][1] = _xor(s[0][1], _mult_gf_24("0100", s[1][1]))
    s_prime[1][1] = _xor(_mult_gf_24("0100", s[0][1]), s[1][1])
    return s_prime[0][0] + s_prime[1][0] + s_prime[0][1] + s_prime[1][1]

def _xor(value1, value2):
    """
    Exclusive-OR
    """
    result = ""
    if len(value1) <= len(value2):
        while len(value1) < len(value2):
            value1 = "0" + value1
    else:
        while len(value2) < len(value1):
            value2 = "0" + value2
    for i in xrange(len(value1)):
        if value1[i] != value2[i]:
            result += "1"
        else:
            result += "0"
    return result

def _mult_gf_24(value1, value2):
    mult_gf = [
               [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
               [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
               [0, 2, 4, 6, 8, 10, 12, 14, 3, 1, 7, 5, 11, 9, 15, 13],
               [0, 3, 6, 5, 12, 15, 10, 9, 11, 8, 13, 14, 7, 4, 1, 2],
               [0, 4, 8, 12, 3, 7, 11, 15, 6, 2, 14, 10, 5, 1, 13, 9],
               [0, 5, 10, 15, 7, 2, 13, 8, 14, 11, 4, 1, 9, 12, 3, 6],
               [0, 6, 12, 10, 11, 13, 7, 1, 5, 3, 9, 15, 14, 8, 2, 4],
               [0, 7, 14, 9, 15, 8, 1, 6, 13, 10, 3, 4, 2, 5, 12, 11],
               [0, 8, 3, 11, 6, 14, 5, 13, 12, 4, 15, 7, 10, 2, 9, 1],
               [0, 9, 1, 8, 2, 11, 3, 10, 4, 13, 5, 12, 6, 15, 7, 14],
               [0, 10, 7, 13, 14, 4, 9, 3, 15, 5, 8, 2, 1, 11, 6, 12],
               [0, 11, 5, 14, 10, 1, 15, 4, 7, 12, 2, 9, 13, 6, 8, 3],
               [0, 12, 11, 7, 5, 9, 14, 2, 10, 6, 1, 13, 15, 3, 4, 8],
               [0, 13, 9, 4, 1, 12, 8, 5, 2, 15, 11, 6, 3, 14, 10, 7],
               [0, 14, 15, 1, 13, 3, 2, 12, 9, 7, 6, 8, 4, 10, 11, 5],
               [0, 15, 13, 2, 9, 6, 4, 11, 1, 14, 12, 3, 8, 7, 5, 10]
              ]
    value1 = int(value1, 2)
    value2 = int(value2, 2)
    return BINARY_VALUE[mult_gf[value1][value2]]

